/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.plugin.okhttp;

import com.lvyh.lightframe.tracer.core.SimpleSpan;
import com.lvyh.lightframe.tracer.core.SimpleTracer;
import com.lvyh.lightframe.tracer.core.common.AbstractTracer;
import com.lvyh.lightframe.tracer.core.common.TraceContants;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContext;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContextHolder;
import com.lvyh.lightframe.tracer.core.samplers.Sampler;
import io.opentracing.propagation.Format;
import okhttp3.Headers;
import okhttp3.Request;
import okhttp3.Response;

import java.io.IOException;
import java.util.List;

public class OkHttpRequestInterceptor implements okhttp3.Interceptor {

    protected String appName;

    protected String targetAppName;

    protected AbstractTracer okHttpTracer;

    public OkHttpRequestInterceptor(AbstractTracer okHttpTracer, String appName, String targetAppName) {
        this.okHttpTracer = okHttpTracer;
        this.appName = appName;
        this.targetAppName = targetAppName;
    }

    @Override
    public Response intercept(Chain chain) throws IOException {
        Request request = chain.request();
        SimpleSpan simpleSpan = okHttpTracer.clientSend(request.method());
        Response response = chain.proceed(setRequestSpanTag(request, simpleSpan));

        if (simpleSpan != null) {
            SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
            SimpleSpan currentSpan = simpleTraceContext.getCurrentSpan();
            Sampler sampler = okHttpTracer.getSimpleTracer().getSampler();
            //sampler is support &  current span is root span
            if (sampler != null && currentSpan.getParentSpan() == null) {
                currentSpan.getSimpleSpanContext().setSampled(true);
                currentSpan.getSimpleSpanContext().setSampler(sampler);
            }
        }

        okHttpTracer.clientReceive(String.valueOf(response.code()));
        return response;
    }

    private Request setRequestSpanTag(Request request, SimpleSpan simpleSpan) {
        String methodName = request.method();
        simpleSpan.setTag(TraceContants.LOCAL_APP, appName == null ? "default" : appName);
        simpleSpan.setTag(TraceContants.REQUEST_URL, request.url().toString());
        simpleSpan.setTag(TraceContants.METHOD, methodName);

        Headers headers = request.headers();
        if (headers != null) {
            List<String> contentLengthList = headers.values("Content-Length");
            if (contentLengthList != null && !contentLengthList.isEmpty()) {
                simpleSpan.setTag(TraceContants.REQ_SIZE, Long.valueOf(contentLengthList.get(0)));
            }
        } else {
            simpleSpan.setTag(TraceContants.REQ_SIZE, String.valueOf(-1));
        }

        return this.injectCarrier(request, simpleSpan);
    }

    public Request injectCarrier(Request request, SimpleSpan currentSpan) {
        Headers.Builder headerBuilder = request.headers().newBuilder();

        SimpleTracer simpleTracer = this.okHttpTracer.getSimpleTracer();
        simpleTracer.inject(currentSpan.getSimpleSpanContext(),
                Format.Builtin.HTTP_HEADERS, new OkHttpRequestCarrier(headerBuilder));

        Request.Builder requestBuilder = request.newBuilder();
        return requestBuilder.headers(headerBuilder.build()).build();
    }
}
